use ::common::compose::{CallBox, Composer, FError};
use ::common::utils::count_tokens;
use text_splitter::{MarkdownSplitter, TextSplitter, CodeSplitter};

pub mod common;
pub mod grap_site;
pub mod to_markdown;
pub mod trans_site;

/// 文本类型枚举
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TextType {
    /// Markdown 文本
    Markdown,
    /// HTML 文本
    Html,
    /// 纯文本
    Txt,
}

/// 文本分割配置
#[derive(Debug, Clone)]
pub struct TextSplitterConfig {
    /// 文本类型
    pub text_type: TextType,
    /// 最大 token 数量
    pub max_tokens: usize,
}

impl Default for TextSplitterConfig {
    fn default() -> Self {
        Self {
            text_type: TextType::Markdown,
            max_tokens: 5000,
        }
    }
}

/// 创建文本分割器 CallBox
///
/// 该 CallBox 接收一个字符串，根据指定的文本类型和最大 token 数量进行分割，
/// 返回分割后的字符串列表。
///
/// # 参数
///
/// * `config` - 文本分割配置，包含文本类型和最大 token 数量
///
/// # 返回值
///
/// 返回一个 CallBox，该 CallBox 接收一个字符串，返回分割后的字符串列表
pub fn make_text_splitter_callbox(config: TextSplitterConfig) -> CallBox<'static, String, Vec<String>> {
    let max_tokens: usize = config.max_tokens;
    let text_type = config.text_type;

    Composer::from_fn(move |text: String| {
        let max_tokens = max_tokens;
        let text_type = text_type;

        async move {
            // 首先使用第三方分割器进行初步分割
            let initial_chunks: Vec<String> = match text_type {
                TextType::Markdown => {
                    // 使用 MarkdownSplitter 进行语义分割
                    let splitter = MarkdownSplitter::new(max_tokens);
                    splitter.chunks(&text).map(|s| s.to_string()).collect()
                },
                TextType::Html => {
                    // 使用 CodeSplitter 配合 tree-sitter-html 进行语义分割
                    let splitter = CodeSplitter::new(
                        tree_sitter_html::LANGUAGE,
                        max_tokens
                    ).map_err(|e| FError::Generic(Box::new(std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))))?;

                    splitter.chunks(&text).map(|s| s.to_string()).collect()
                },
                TextType::Txt => {
                    // 使用基本的 TextSplitter 进行分割
                    let splitter = TextSplitter::new(max_tokens);
                    splitter.chunks(&text).map(|s| s.to_string()).collect()
                },
            };

            // 如果初步分割结果为空或只有一个块，直接返回
            if initial_chunks.is_empty() || initial_chunks.len() == 1 {
                return Ok(initial_chunks);
            }

            // 对初步分割的块进行二次处理，尽可能合并小块
            let mut optimized_chunks: Vec<String> = Vec::new();
            let mut current_chunk = String::new();
            let mut current_tokens = 0;

            // 计算目标 token 数量，使用最大值的90%作为目标，避免过度接近最大值
            let target_tokens = (max_tokens as f64 * 0.9) as usize;

            for chunk in initial_chunks {
                // 计算当前块的 token 数量
                let chunk_tokens = count_tokens(&chunk).unwrap_or(chunk.len() / 4); // 如果计算失败，使用估计值

                // 如果当前块加上新块不会超过目标 token 数量，则合并
                if current_tokens + chunk_tokens <= target_tokens {
                    // 如果当前块不为空，添加分隔符
                    if !current_chunk.is_empty() {
                        current_chunk.push_str("\n\n");
                    }
                    current_chunk.push_str(&chunk);
                    current_tokens += chunk_tokens;
                } else {
                    // 如果当前块不为空，将其添加到结果中
                    if !current_chunk.is_empty() {
                        optimized_chunks.push(current_chunk);
                        current_chunk = chunk;
                        current_tokens = chunk_tokens;
                    } else {
                        // 如果当前块为空，说明新块本身就超过了目标 token 数量
                        optimized_chunks.push(chunk);
                        current_chunk = String::new();
                        current_tokens = 0;
                    }
                }
            }

            // 处理最后一个块
            if !current_chunk.is_empty() {
                optimized_chunks.push(current_chunk);
            }

            Ok(optimized_chunks)
        }
    })
}